﻿#include "Screenshot.h"
#include "tchar.h"
#include "dwmapi.h"
#include <QDebug>

#pragma comment(lib,"dwmapi.lib")

DibCaptureHelper::DibCaptureHelper()
    : hwnd_(nullptr)
    , scrDc_(nullptr)
    , memDc_(nullptr)
    , bitmap_(nullptr)
    , oldBitmap_(nullptr)
    , windowRect_{ 0, 0, 0, 0 }
    , clientRect_{ 0, 0, 0, 0 }
    , bmpDataSize_(0)
{
}

DibCaptureHelper::~DibCaptureHelper()
{
    Cleanup();
}

bool DibCaptureHelper::Init(HWND hwnd)
{
    hwnd_ = hwnd;

    // 获取窗口大小
    if (!::GetWindowRect((HWND)_capturer.hwnd, &windowRect_) || 
        !::GetClientRect((HWND)_capturer.hwnd, &clientRect_))
    {
        return false;
    }

    if (_capturer.hwnd != GetDesktopWindow())
    {
        DwmGetWindowAttribute((HWND)_capturer.hwnd, DWMWA_EXTENDED_FRAME_BOUNDS, (PVOID)&clientRect_, sizeof(clientRect_));
    }

    const auto clientRectWidth = clientRect_.right - clientRect_.left;
    const auto clientRectHeight = clientRect_.bottom - clientRect_.top;

    // 位图信息
    BITMAPINFO bitmapInfo;
    bitmapInfo.bmiHeader.biSize = sizeof(BITMAPINFOHEADER);
    bitmapInfo.bmiHeader.biBitCount = 24;
    bitmapInfo.bmiHeader.biCompression = BI_RGB;
    bitmapInfo.bmiHeader.biWidth = clientRectWidth;
    bitmapInfo.bmiHeader.biHeight = clientRectHeight;
    bitmapInfo.bmiHeader.biPlanes = 1;
    bitmapInfo.bmiHeader.biSizeImage = clientRectWidth * clientRectHeight * bitmapInfo.bmiHeader.biBitCount / 8;

    bmpDataSize_ = bitmapInfo.bmiHeader.biSizeImage;

    scrDc_ = GetWindowDC(hwnd_); // 获取窗口DC
    memDc_ = CreateCompatibleDC(scrDc_); // 缓冲内存DC
    bitmap_ = CreateDIBSection(scrDc_, &bitmapInfo, DIB_RGB_COLORS, &_capturer.buffer, nullptr, 0);

    if (bitmap_ == nullptr)
    {
        DeleteDC(memDc_);
        ReleaseDC(hwnd_, scrDc_);
        return false;
    }

    oldBitmap_ = static_cast<HBITMAP>(SelectObject(memDc_, bitmap_));

    return true;
}

void DibCaptureHelper::Cleanup()
{
    if (bitmap_ == nullptr)
    {
        return;
    }

    // 删除用过的对象
    SelectObject(memDc_, oldBitmap_);

    if (!DeleteObject(bitmap_))
    {
        printf("\n DeleteObject faild\n");
    }

    if (!DeleteDC(memDc_))
    {
        printf("\n DeleteDC faild\n");
    }

    ReleaseDC(hwnd_, scrDc_);

    hwnd_ = nullptr;
    scrDc_ = nullptr;
    memDc_ = nullptr;
    bitmap_ = nullptr;
    oldBitmap_ = nullptr;
}

bool DibCaptureHelper::RefreshWindow()
{
    const auto hwnd = hwnd_;
    Cleanup();
    return Init(hwnd);
}

bool DibCaptureHelper::ChangeWindowHandle(const Capturer& capturer)
{
    _capturer = capturer;

    Cleanup();

    if (capturer.type == 1)
    {
        return Init(GetDesktopWindow());
    }
    else
    {
        return Init((HWND)capturer.hwnd);
    }
}

bool DibCaptureHelper::Capture()
{
    if (bitmap_ == nullptr || memDc_ == nullptr || scrDc_ == nullptr)
    {
        return false;
    }

    if (!IsWindow((HWND)_capturer.hwnd))
    {
        return false;
    }

    // 获取窗口大小
    if (_capturer.hwnd != GetDesktopWindow())
    {
        RECT capturerRect;
        DwmGetWindowAttribute((HWND)_capturer.hwnd, DWMWA_EXTENDED_FRAME_BOUNDS, (PVOID)&capturerRect, sizeof(capturerRect));

        if (capturerRect.left != clientRect_.left || capturerRect.right != clientRect_.right
            || capturerRect.bottom != clientRect_.bottom || capturerRect.top != clientRect_.top)
        {
            RefreshWindow();
        }

        _capturer.width = capturerRect.right - capturerRect.left;
        _capturer.height = capturerRect.bottom - capturerRect.top;

        _capturer.x = capturerRect.left;
        _capturer.y = capturerRect.top;
    }

    bool result = false;

    switch (_capturer.type)
    {
    case 0: // 句柄截图
        if (_capturer.hwnd == GetDesktopWindow())
        {
            result = BitBlt(memDc_, 0, 0, _capturer.width, _capturer.height,
                scrDc_, _capturer.x, _capturer.y, SRCCOPY);
        }
        else
        {
            result = BitBlt(memDc_, 0, 0, _capturer.width, _capturer.height,
                scrDc_, 0, 0, SRCCOPY);
        }
        break;
    case 1: // 区域截图
        if (_capturer.hwnd != GetForegroundWindow() 
            && _capturer.hwnd != GetDesktopWindow())
        {
            return true;
        }
        result = BitBlt(memDc_, 0, 0, _capturer.width, _capturer.height,
            scrDc_, _capturer.x, _capturer.y, SRCCOPY);
        break;
    case 2: // 句柄截图
        result = PrintWindow((HWND)_capturer.hwnd, memDc_, PW_CLIENTONLY);
        break;
    default:
        break;
    }

    /* 绘制鼠标 */
    if (_capturer.cursor)
    {
        CURSORINFO cursorInfo;
        cursorInfo.cbSize = sizeof(CURSORINFO);
        GetCursorInfo(&cursorInfo);
        POINT pt = cursorInfo.ptScreenPos;
        ScreenToClient((HWND)_capturer.hwnd, &pt);

        if (_capturer.hwnd == GetDesktopWindow())
        {
            pt.x -= _capturer.x;
            pt.y -= _capturer.y;
        }

        DrawIcon(memDc_, pt.x, pt.y, cursorInfo.hCursor);
    }

    return result;
}

void DibCaptureHelper::EnumCapturers(std::vector<Capturer>* capturers)
{
    EnumDisplayMonitors(nullptr, nullptr, MonitorEnumProc, (LPARAM)capturers);
    EnumWindows(EnumWindowsProc, (LPARAM)capturers);
}

BOOL DibCaptureHelper::MonitorEnumProc(HMONITOR hMonitor, HDC hdcMonitor, LPRECT lprcMonitor, LPARAM dwData)
{
    std::vector<Capturer>* vec = (std::vector<Capturer>*)dwData;

    Capturer cap;

    cap.x = lprcMonitor->left;
    cap.y = lprcMonitor->top;
    cap.width = lprcMonitor->right - lprcMonitor->left;
    cap.height = lprcMonitor->bottom - lprcMonitor->top;

    MONITORINFOEXA mi;
    mi.cbSize = sizeof(mi);
    GetMonitorInfoA(hMonitor, &mi);

    char lastChar = mi.szDevice[strlen(mi.szDevice) - 1];

    cap.hwnd = GetDesktopWindow();
    cap.name = lastChar;

    vec->push_back(cap);
    return true;
}

BOOL DibCaptureHelper::EnumWindowsProc(HWND hwnd, LPARAM lParam)
{
    std::vector<Capturer>* vec = (std::vector<Capturer>*)lParam;

    char szClass[256] = { 0 };
    int nMaxCount = 256;

    LPSTR lpClassName = szClass;
    GetClassNameA(hwnd, lpClassName, nMaxCount);

    if (strcmp(lpClassName, "Windows.UI.Core.CoreWindow") == 0 
        || strcmp(lpClassName, "ApplicationFrameWindow") == 0 
        || strcmp(lpClassName, "Progman") == 0
        || strcmp(lpClassName, "Shell_TrayWnd") == 0
        || strcmp(lpClassName, "CEF-OSC-WIDGET") == 0)
    {
        return true;
    }

    // 窗口是否可视
    if (!IsWindowVisible(hwnd))
    {
        return true;
    }

    // 窗口是否可激活
    if (!IsWindowEnabled(hwnd))
    {
        return true;
    }

    // 过滤自己的窗口
    DWORD wndPid = 0;
    DWORD curPid = GetCurrentProcessId();
    GetWindowThreadProcessId(hwnd, &wndPid);

    if (wndPid == curPid)
    {
        return true;
    }

    RECT rc = { 0 };
    ::GetWindowRect(hwnd, &rc);

    Capturer cap;

    cap.x = rc.left;
    cap.y = rc.top;
    cap.width = rc.right - rc.left;
    cap.height = rc.bottom - rc.top;

    char title[128] = { 0 };
    GetWindowTextA(hwnd, title, 128);

    cap.name = title;
    cap.hwnd = hwnd;

    if (!cap.name.empty())
    {
        vec->push_back(cap);
    }

    return true;
}
